
import tensorflow as tf
import glob
import shutil
import xml.etree.ElementTree as ET
import os

def pix1280to512(path):
    imgs = glob.glob(path + "*.jpg")
    print(imgs)
    for i in imgs:
        # print(i)
        img = tf.io.read_file(i)
        img = tf.image.decode_jpeg(img, channels=3)
        img = tf.image.resize_with_pad(img, 512, 512, # 'nearest',
                                           antialias=True)
        img = tf.clip_by_value(img, 0, 255)
        img = tf.cast(img, dtype=tf.uint8)

        img = tf.image.encode_jpeg(img)

        with tf.io.gfile.GFile(i, 'wb') as file:
            file.write(img.numpy())


def pix1280to512_xml(path):
    xmls = glob.glob(path + "*.xml")
    print(xmls)
    for i in xmls:
        print(i)
        tree = ET.parse(i)

        img_info = dict()
        boxes_counter = 0
        img_info['object'] = []
        for elem in tree.iter():

            if 'object' in elem.tag or 'part' in elem.tag:
                for attr in list(elem):
                    if 'bndbox' in attr.tag:
                        for pos in list(attr):
                            if 'xmin' in pos.tag:
                                # object_info[0] = int(pos.text)
                                pos.text = str(int(float(pos.text) / 2.5))
                                # object_info[0] = float(pos.text) / 2.5
                            if 'ymin' in pos.tag:
                                # object_info[1] = int(pos.text)
                                pos.text = str(int((float(pos.text) + 280.) / 2.5))
                                # object_info[1] = (float(pos.text) + 280.) / 2.5
                            if 'xmax' in pos.tag:
                                # object_info[2] = int(pos.text)
                                pos.text = str(int(float(pos.text) / 2.5))
                                # object_info[2] = float(pos.text) / 2.5
                            if 'ymax' in pos.tag:
                                # object_info[3] = int(pos.text)
                                pos.text = str(int((float(pos.text) + 280.) / 2.5))
                                # object_info[3] = (float(pos.text) + 280.) / 2.5
        tree.write(i)


def getTestImg():
    allImgPath = glob.glob("data/train_512/*.jpg")
    allImgPath = [int(i[15:-4]) for i in allImgPath]
    print(allImgPath)
    trainImgPath = glob.glob("data/train2/*.jpg")
    trainImgPath = [int(i[12:-4]) for i in trainImgPath]
    print(trainImgPath)

    allXmlPath = glob.glob("data/trainxml/*.xml")
    allXmlPath = [int(i[14:-4]) for i in allXmlPath]
    print(allXmlPath)
    trainXmlPath = glob.glob("data/train2xml/*.xml")
    trainXmlPath = [int(i[15:-4]) for i in trainXmlPath]
    print(trainXmlPath)

    assert allImgPath.__len__() == allXmlPath.__len__()
    assert trainImgPath.__len__() == trainXmlPath.__len__()


    for i in allImgPath:
        if not i in trainImgPath:

            source = "data\\train_512\\" + str(i) + ".jpg"
            target = "data\\test2\\" +  str(i) + ".jpg"

            try:
                shutil.copy(source, target)
            except IOError as e:
                print("Unable to copy file. %s" % e)


    for i in allXmlPath:
        if not (i in trainXmlPath):
            source = "data\\trainxml\\" + str(i) + ".xml"
            target = "data\\test2xml\\" + str(i) + ".xml"

            try:
                shutil.copy(source, target)
            except IOError as e:
                print("Unable to copy file. %s" % e)


def deviceData():
    imgsPath = "data\\___train\\"
    labelsPath = "data\\___trainxml\\"
    trueImgsName = glob.glob(imgsPath + "*.jpg")
    trueLabelsName = glob.glob(labelsPath + "*.xml")


    for i in range(3500):
        img = str(i) + ".jpg"
        label = str(i) + ".xml"

        if (1580 <= i <= 2139) or (2390 <= i <= 2613):
            n = 2
        else :
            n = 4
        if i % n == 0:
            imgSavePath = "data\\trainDarkNet53\\"
            labelSavePath = "data\\trainDarkNet53Xml\\"
        else:
            imgSavePath = "data\\testDarkNet53\\"
            labelSavePath = "data\\testDarkNet53Xml\\"

        if ((imgsPath + img) in trueImgsName) and ((labelsPath + label) in trueLabelsName):
            try:
                shutil.copy(imgsPath + img, imgSavePath + img)
                shutil.copy(labelsPath + label, labelSavePath + label)
            except IOError as e:
                print("Unable to copy file. %s" % e)



def deviceTestAndValData(path):

    testAndValImgPath = path + "testandvalimg\\"
    testAndValLabelPath = path + "testandvallabel\\"
    testImgPath = path + "testimg\\"
    testLabelPath = path + "testlabel\\"
    valImgPath = path + "valimg\\"
    valLabelPath = path + "vallabel\\"

    beginIdx = 0
    endIdx = 5000

    deviceRecord = 0

    for idx in range(beginIdx, endIdx + 1):
        testAndValImg = testAndValImgPath + str(idx) + ".jpg"
        testAndValLabel = testAndValLabelPath + str(idx) + ".xml"
        testImg = testImgPath + str(idx) + ".jpg"
        testLabel = testLabelPath + str(idx) + ".xml"
        valImg = valImgPath + str(idx) + ".jpg"
        valLabel = valLabelPath + str(idx) + ".xml"


        if deviceRecord == 3:
        # if idx % 4 == 0 or idx % 8 == 0:
            try:
                shutil.copy(testAndValImg, testImg)
                shutil.copy(testAndValLabel, testLabel)
                if os.path.exists(testAndValImg) and os.path.exists(testAndValLabel):
                    deviceRecord = 0
            except IOError as e:
                print("Unable to copy file. %s" % e)
        else:
            try:
                shutil.copy(testAndValImg, valImg)
                shutil.copy(testAndValLabel, valLabel)
                if os.path.exists(testAndValImg) and os.path.exists(testAndValLabel):
                    deviceRecord += 1
            except IOError as e:
                print("Unable to copy file. %s" % e)




def copySameXml(needCopyImgPath, needCopyXmlPath, copyedXmlPath):
    il = glob.glob(needCopyImgPath + "*.jpg")
    il = [i.split("\\")[-1][:-3] + "xml" for i in il]
    for i in il:
        shutil.copy(copyedXmlPath + i, needCopyXmlPath + i)



def xmlRemoveSmallBox(xmlPath, savePath, SMALL_SIZE=5000):
    tree = ET.parse(xmlPath)
    root = tree.getroot()
    removeIters = []
    for elem in tree.iter():
        if 'object' in elem.tag:
            removeFlag = False
            for attr in list(elem):
                if 'bndbox' in attr.tag:
                    for pos in list(attr):
                        if 'xmin' in pos.tag:
                            xmin = int(pos.text)
                        if 'ymin' in pos.tag:
                            ymin = int(pos.text)
                        if 'xmax' in pos.tag:
                            xmax = int(pos.text)
                        if 'ymax' in pos.tag:
                            ymax = int(pos.text)
                    assert xmin != None
                    assert ymin != None
                    assert xmax != None
                    assert ymax != None
                    if (xmax - xmin) * (ymax - ymin) < SMALL_SIZE:
                        removeFlag = True
            if removeFlag:
                removeIters.append(elem)
    for i in removeIters:
        root.remove(i)

    # with open(savePath, "w", encoding="UTF-8") as f:
    #     tree.writexml(f, addindent='  ', encoding='utf-8')
    tree.write(savePath)
    # tree.write(savePath, encoding='utf-8')


if __name__ == '__main__':
    # getTestImg()
    # deviceData()
    # deviceTestAndValData("newData\\")
    # copySameXml("data\\valDarkNet53\\", "data\\valDarkNet53Xml\\", "data\\testDarkNet53xml\\")
    # pix1280to512("1\\")

    pix1280to512("data20230208\\testimgs\\")
    # pix1280to512_xml("data20230208\\testlabels\\")

    pix1280to512("data20230208\\trainimgs\\")
    pix1280to512_xml("data20230208\\trainlabels\\")

    # pix1280to512("data\\zhuaqu\\*\\")
    # pix1280to512("newData\\trainimg\\")
    # pix1280to512_xml("newData\\trainlabel\\")
    # pix1280to512("newData\\testimg\\")
    # pix1280to512_xml("newData\\testlabel\\")
    # pix1280to512("newData\\valimg\\")
    # pix1280to512_xml("newData\\vallabel\\")

    # primerXmls = glob.glob("data\\___trainxml\\*.xml")
    # allSavePath = "data\\___newTrainXml\\"
    #
    # for i in primerXmls:
    #     savePath = allSavePath + i.split("\\")[-1]
    #     xmlRemoveSmallBox(i, savePath, 5000)

    # il = glob.glob("data\\___newTrain\\*.jpg")
    # xl = glob.glob("data\\___newTrainXml\\*.xml")
    # il = [int(i.split("\\")[-1][:-4]) for i in il]
    # xl = [int(i.split("\\")[-1][:-4]) for i in xl]
    # for i in range(1500, 3000):
    #     # istr = str(i)
    #     if ((not i in il) and ( i in xl)):
    #         print(i)